package TestPackage;

import DataSetPackage.DataSet;
import DataSetPackage.DataSetStorage;
import DataSetPackage.ProductView;
import DataSetPackage.SimpleView;
import DataSetPackage.View;
import EstimatorPackage.EmpSuffix;
import EstimatorPackage.EstimatorExp;
import EstimatorPackage.IndependentProductEst;
import EstimatorPackage.LaPlaceSuffix;
import EstimatorPackage.SimSuffixEstimatorExp;
import HierarchyPackage.Hierarchy;
import HierarchyPackage.ProductHierarchy.ProductHierarchy;
import HierarchyPackage.AffixHierarchies.SuffixHierarchy;
import HierarchyPackage.AffixHierarchies.WordHierarchy;
import HierarchyPackage.HClass;
import HierarchyPackage.ProductHierarchy.ProductHClass;
import HierarchyPackage.SimpleHierarchy.SimpleHierarchy;
import HierarchyPackage.SimpleHierarchy.SimpleHierarchyMem;
import OptionsManager.DiscountOptions;
import OptionsManager.TrainOptions;
import OrderedCoverPackage.BaseHPM;
import OrderedCoverPackage.NSuffixOCEstimator;
import OrderedCoverPackage.NWordOCEstimator;
import OrderedCoverPackage.OrderedCoverEstimator;
import java.io.BufferedReader;
import java.io.FileReader;
import java.io.FileWriter;
import java.io.PrintWriter;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import org.osdtsystem.utils.MemoryStorage;
import org.osdtsystem.utils.Storage;
import perplexity.EstimatorPer;
import perplexity.OCsufestimator;
import perplexity.OCwordestimator;
import perplexity.PerplexityCalculator;
import util.IO;
import util.WordPOS;

/**
 *
 * @author Martin Haulrich
 */
public class Main {

    static String newline = System.getProperty("line.separator");

    public static final String SOSPOS = "<SOS>";

    public static final String EOSPOS = "<EOS>";

    /** Constructor */
    public Main() {
    }

    /**
     * @param args the command line arguments
     */
    public static void main(String[] args) throws Exception {

        if (false) {
            SimpleHierarchy POSHier = new SimpleHierarchyMem("output/SUC.hier");
            System.err.println(POSHier);

            Storage storage = new MemoryStorage();
            DataSet ds = new DataSetStorage(storage);

            List<List<WordPOS>> train = IO.readSUC("/Users/mwh/Java/HMM-Tagger/train.txt");

            List<String> priorTrain = new ArrayList<String>();

            for (List<WordPOS> sen : train) {
                List<String> dp = new ArrayList<String>(2);
                dp.add(SOSPOS);
                dp.add(sen.get(0).POS);
                ds.add(dp);
                for (int i = 1; i < sen.size(); i++) {
                    dp = new ArrayList<String>(2);
                    dp.add(sen.get(i - 1).POS);
                    dp.add(sen.get(i).POS);
                    ds.add(dp);
                }
                dp = new ArrayList<String>(2);
                dp.add(sen.get(sen.size() - 1).POS);
                dp.add(EOSPOS);
                ds.add(dp);

                priorTrain.add(SOSPOS);
                for (int i = 0; i < sen.size(); i++) {
                    priorTrain.add(sen.get(i).POS);
                }
                priorTrain.add(EOSPOS);
            }

            EmpPOS es = new EmpPOS(priorTrain);

            EmpPOS ls = new LaPlacePOS(es);
            // Create estimator to handle expression
            EstimatorExp ees = new SimPOSEstimatorExp(es, POSHier);
            EstimatorExp els = new SimPOSEstimatorExp(ls, POSHier);

            EstimatorExp[] ee = {els, ees};

            // Create product-estimator
            IndependentProductEst ie = new IndependentProductEst(ee);

            EstimatorExp pe = ie;

            System.err.println("DatasetSize: " + ds.size());

            // Data for hierarchy is at index 0.

            SimpleView posdv1 = new SimpleView(0, ds);
            SimpleView posdv2 = new SimpleView(1, ds);

            // Create productView
//            View[] views = {sufdv1, sufdv2};
            View[] views = {posdv1, posdv2};
            ProductView pv = new ProductView(ds, views);


            System.err.println(ds.getFeatures(4));
            System.err.println(es.getEstimate("PP"));
            System.err.println(ees.getEstimate(4, posdv1));
            System.err.println(ees.getEstimate(4, posdv2));
            System.err.println(ie.getEstimate(4, pv));


            for (Integer did : ds.getIDs()) {
//                System.err.println(ds.getFeatures(did) + "\t" + pe.getEstimate(did, pv));
            }

//            Hierarchy[] hiers = {sufhier, sufhier};
            Hierarchy[] hiers = {POSHier, POSHier};

            // And now we make the suf x suf producthierarchy
            ProductHierarchy phier = new ProductHierarchy(hiers);

            System.err.println("Creating cover...");
            // BaseHPM(hierarchy, prior, mincount, searchBreadth)

            TrainOptions options = new TrainOptions();
            DiscountOptions discountOptions = new DiscountOptions();

            BaseHPM oc = new BaseHPM(phier, pe, discountOptions);
            if (true) {
                System.err.println("Training cover...\n");

                long a = System.currentTimeMillis();
                oc.train(ds.getIDs(), pv, options);
                long b = System.currentTimeMillis();
                long t = b - a;

                System.err.println("Training time: " + t);
//
                oc.toDot2("/Users/mwh/Java/XHPM-MWH/output/pos.dot", false);
//

                System.err.println("CoverSize: " + oc.OCsize());
                System.err.println("BumpPost: " + oc.getBumpPostMDL(pv));


                OrderedCoverEstimator oce = new OrderedCoverEstimator(oc);

                for (int i : POSHier.getAllClasses()) {


                    HClass c1 = POSHier.getHClass(i);
                    if (!c1.equals(POSHier.getTop())) {
                        double sum = 0;
//                        System.err.println(c1.getName());
                        for (int j : POSHier.getAllClasses()) {
                            HClass c2 = POSHier.getHClass(j);
                            if (!c2.equals(POSHier.getTop())) {
                                System.err.print(c1.getName() + "\t" + c2.getName());
                                HClass[] classes = {c1, c2};
                                ProductHClass pc = new ProductHClass(phier, classes);
                                double cond = oce.getConditioningProb(pc);
                                double prob = oce.classAsDataPointEstimate(pc) - cond;
                                System.err.println("\t" + prob);
                                sum = sum + Math.pow(10, prob);
                            }
                        }
                        System.err.println("Sum: " + sum);
                    }
                }
//                NSuffixOCEstimator oce = new NSuffixOCEstimator(oc);

//                oce.save(args[2]);
//                PrintWriter out = new PrintWriter(new FileWriter(args[2] + ".log"));
//                out.println("Data points: " + oc.dataSize());
//                out.println("Train options: " + options);
//                out.println("");
//                out.println("Training time:" + t + "ms");
//                out.println("Cover size: " + oc.OCsize());
//                out.println("Score: " + oc.getBumpPostMDL(pv));
//                out.close();
            }

        }

        // Arguments:
        // 1. options file
        // 2. discount options file
        // 3. language
        // 4. N
        // 5. Train data file
        // 6. train data size
        // 7. evalutation file
        // 8. evaulate goodness-of-fit (yes/no)

        if (true) {

            long tta = System.currentTimeMillis();
            TrainOptions options = new TrainOptions(args[0]);
            DiscountOptions discountOptions = new DiscountOptions(args[1]);
//            options.setIntOption("mincount", 2);
//            options.setIntOption("depth", 10);
//            options.setIntOption("rounds", 2);
            System.err.println(options);

            String language = args[2];
            int N = Integer.parseInt(args[3]);
            String trainDataSize = args[5];

            String evalFile = args[6];

            Storage storage = new MemoryStorage();
            String fn = args[4];
            DataSet ds = new DataSetStorage(storage);

            boolean evalTrain = false;
            if (args[7].equals("yes")) {
                evalTrain = true;

            }

            BufferedReader input = new BufferedReader(new FileReader(fn));

            // Just a counter to control how much data is readb
            int rl = 0;

            // List of words. Used to create empirical estimator
            List<String> wl = new ArrayList<String>();

            int maxread = Integer.MAX_VALUE;//Integer.parseInt(args[1]);

            int lc = 0;
            String line;
            System.err.println("Reading data...");

//            List<List<String>> ngrams = new ArrayList<List<String>>();
            while (((line = input.readLine()) != null) && rl < maxread) {
                rl++;
                for (int i = 0; i < N - 1; i++) {
                    line = "<s> " + line;
                }
                line = line + " </s>";
                String[] tokens = line.split(" ");
//                System.err.println(line);
//                for (int i = 0; i < tokens.length - 1; i++) {
//                    List<String> ng = new ArrayList<String>(2);
//                    ng.add("^" + tokens[i]);
//                    ng.add("^" + tokens[i + 1]);
//                    ngrams.add(ng);
//                    ds.add(ng);
//             log2   }

                for (int i = 0; i < tokens.length - (N - 1); i++) {
                    List<String> ng = new ArrayList<String>(N);
                    for (int j = 0; j < N; j++) {
                        ng.add("^" + tokens[i + j]);
                    }

                    ds.add(ng);
//                    System.err.println(ng);
                }
//                for (String t : tokens) {
                for (int i = (N - 2); i < tokens.length; i++) {

                    String t = tokens[i];
                    // Add prefix to word
                    String w = "^" + t;
                    wl.add(w);
                }

                lc++;
            }
            input.close();

            Set<Character> alphabet = new HashSet<Character>();
            for (String w : wl) {
                for (char c : w.toCharArray()) {
                    alphabet.add(c);
                }
            }


            System.err.println("Lines read: " + lc);
            System.err.println("words read: " + wl.size());
            System.err.println("");

            System.err.println("Creating prior...");
            EmpSuffix es = new EmpSuffix(wl);

            EmpSuffix ls = new LaPlaceSuffix(es);
            // Create estimator to handle expression
            EstimatorExp ees = new SimSuffixEstimatorExp(es);
            EstimatorExp els = new SimSuffixEstimatorExp(ls);



            // Array of estimators used in product-estimator
            //EstimatorExp[] ee = {se, se};
//            EstimatorExp[] ee = {els, ees};
            EstimatorExp[] ee = new EstimatorExp[N];
            for (int i = 0; i < (N - 1); i++) {
                ee[i] = els;
            }
//            ee[N - 1] = els;
            ee[N - 1] = ees;

            // Create product-estimator
            IndependentProductEst ie = new IndependentProductEst(ee);

            EstimatorExp pe = ie;

            System.err.println("Creating product dataset");
            // Make empty dataset

            System.err.println("DatasetSize: " + ds.size());

            // Create suffix-hiearchy
            SuffixHierarchy sufhier = new SuffixHierarchy();


            View[] views = new View[N];
            for (int i = 0; i < N; i++) {
                views[i] = new SimpleView(i, ds);
            }
            ProductView pv = new ProductView(ds, views);

//            Hierarchy[] hiers = {sufhier, sufhier};

            WordHierarchy whier = new WordHierarchy();

            Hierarchy[] hiers = new Hierarchy[N];
            for (int i = 0; i < N; i++) {
                hiers[i] = sufhier;
//                hiers[i] = whier;
            }


            // And now we make the suf x suf producthierarchy
            ProductHierarchy phier = new ProductHierarchy(hiers);


            // Create OC on the base of the hierarchy, the estimator and a minimun count
            System.err.println("Creating cover...");
            // BaseHPM(hierarchy, prior, mincount, searchBreadth)
            BaseHPM oc = new BaseHPM(phier, pe, discountOptions);

            System.err.println("Training cover...\n");

            long a = System.currentTimeMillis();
            oc.train(ds.getIDs(), pv, options);
            long b = System.currentTimeMillis();
            long t = b - a;

//            StringBuilder sb = new StringBuilder("output/");
  StringBuilder sb = new StringBuilder();
            sb.append(language + "_").append(trainDataSize).append("_N-" + N).append("_D-" + options.getProperty("depth")).
                    append("_M-" + options.getProperty("mincount")).append("_B-" + options.getProperty("breadth")).
                    append("_C-" + options.getProperty("cutfraction")).append("_R-" + options.getProperty("rounds")).
                    append("_d1-" + discountOptions.getProperty("d1")).append("_d2-" + discountOptions.getProperty("d2")).
                    append("_d3-" + discountOptions.getProperty("d3+")).
                    append("_MDL-" + discountOptions.getProperty("MDL"));



            System.err.println("Training time: " + t);

            oc.toDot2(sb.toString() + ".dot", false);


            System.err.println("CoverSize: " + oc.OCsize());
            System.err.println("BumpPost: " + oc.getBumpPostMDL(pv));

            if (false) {
                throw new RuntimeException();
            }

//            NWordOCEstimator oce = new NWordOCEstimator(oc);
            NSuffixOCEstimator oce = new NSuffixOCEstimator(oc);

//
            if (true) {
                double e3 = oce.getEstimate(phier.getTop());
                System.err.println("OC: " + e3 + "\t" + Math.pow(10, e3));
            }




            oce.save(sb.toString());

            System.err.println("HHH 111");
            PrintWriter out = new PrintWriter(new FileWriter(sb.toString() + ".log"));
            out.println("Data points: " + oc.dataSize());
            out.println("Train options: " + options);
            out.println("Dicounts options: " + discountOptions);
            out.println("");
            out.println("Training time: " + t + "ms");
            out.println("Cover size: " + oc.OCsize());
            out.println("Train data likelihood: " + oc.getBumpPostMDL(pv));
            out.println("");
            out.println("--------------");
            out.println("");

            System.err.println("Training time: " + t);

            if (false) {
                throw new RuntimeException();
            }

            EstimatorPer ocefile = new OCsufestimator(fn, sb.toString(), N, pe);
//            EstimatorPer ocefile = new OCwordestimator(fn, sb.toString(), N, pe);

            if (ocefile instanceof OCsufestimator) {
                out.println("");
                out.println("--------------------------------");
                out.println("----- SANITY CHECK -------------");
                out.println("");


                List<String> ng = new ArrayList<String>(N);
                for (int i = 0; i < N - 1; i++) {
                    ng.add("e");
                }
                ng.add("");

                double e = ocefile.getEstimate2(ng);
                out.println("*, \"\": " + e + "\t" + Math.pow(10, e));



                double s = Double.NEGATIVE_INFINITY;
                double s1 = Double.NEGATIVE_INFINITY;
                double s2 = Double.NEGATIVE_INFINITY;


                for (Character c : alphabet) {

//                System.err.println("-------------------------");

                    ng.set((N - 1), c.toString());
//                System.err.println(ng);
                    e = ocefile.getEstimate2(ng);
                    double l = Math.max(s, e);
//                System.err.println(c+" " + e+ "\t" + Math.pow(10, e));
                    s = l + Math.log10(Math.pow(10, s - l) + Math.pow(10, e - l));



                    e = es.getEstimate(c.toString());
                    l = Math.max(s1, e);
//                System.err.println(c+" " + e+ ";" + Math.pow(10, e));
//                System.err.println(Math.pow(10, e));
                    s1 = l + Math.log10(Math.pow(10, s1 - l) + Math.pow(10, e - l));

                    e = ls.getEstimate(c.toString());
                    l = Math.max(s2, e);
//                System.err.println(c+" " + e+ "\t" + Math.pow(10, e));
                    s2 = l + Math.log10(Math.pow(10, s2 - l) + Math.pow(10, e - l));



                }
                ng.set(N - 1, "*");
                out.println("P(" + ng + ") = " + Math.pow(10, s) + "\t\tShould be 1");
                out.println("emp F*(*) = " + Math.pow(10, s1) + "\t\tShould be 1");
                out.println("lap F*(*) = " + Math.pow(10, s2) + "\t\tShould be 1");
                out.println("");


                out.println("--------------------------------");
                out.println("");

                System.err.println("HHHH 222");
            }


            PerplexityCalculator poc = new PerplexityCalculator(ocefile, N);
            double[] res;
            if (evalTrain) {
                a = System.currentTimeMillis();
                res = poc.calculate(fn);
                b = System.currentTimeMillis();
                System.err.println("Calculation took: " + (b - a) + "ms");

                out.println("FIT:");
                out.println("Evaluation time: " + (b - a));
                out.println("unknown: " + res[3]);

                out.println("");
                out.println("logprob: " + res[0]);
                out.println("ppl: " + res[1]);
                out.println("ppl1: " + res[2]);

                out.println("");
                out.println("");
            }

            a = System.currentTimeMillis();
            res = poc.calculate(evalFile);
            b = System.currentTimeMillis();
            System.err.println("Calculation took: " + (b - a) + "ms");

            out.println("PREDICTION:");
            out.println("Evaluation time: " + (b - a));
            out.println("unknown: " + res[3]);

            out.println("");
            out.println("logprob: " + res[0]);
            out.println("ppl: " + res[1]);
            out.println("ppl1: " + res[2]);





            out.close();
            long ttb = System.currentTimeMillis();
            System.err.println("Total time: " + (ttb - tta) + "ms");
        }

    }
}